#include "lockfree_skiplist.h"
#include "../atomic/cpu_pause.h"
#include <stdlib.h>
#include <time.h>

// 生成随机层数
static int random_level(void) {
    int level = 1;
    while ((rand() / (double)RAND_MAX) < PROBABILITY && level < MAX_LEVEL) {
        level++;
    }
    return level;
}

// 创建新节点
static skiplist_node_t* create_node(int key, void* value, int level) {
    skiplist_node_t* node = malloc(sizeof(skiplist_node_t) + 
                                 level * sizeof(atomic_tagged_ptr_t));
    if (!node) return NULL;
    
    node->key = key;
    node->value = value;
    node->level = level;
    
    for (int i = 0; i < level; i++) {
        tagged_ptr_t initial = make_tagged_ptr(NULL, 0);
        atomic_tagged_store(&node->next[i], initial, memory_order_relaxed);
    }
    
    return node;
}

bool skiplist_init(skiplist_t* list, 
                  int (*compare)(const void* key1, const void* key2)) {
    list->head = create_node(INT_MIN, NULL, MAX_LEVEL);
    if (!list->head) return false;
    
    atomic_store(&list->level, 1, memory_order_relaxed);
    list->compare = compare;
    srand(time(NULL));
    
    return true;
}

void skiplist_destroy(skiplist_t* list) {
    skiplist_node_t* current = list->head;
    while (current) {
        skiplist_node_t* next = get_ptr(
            atomic_tagged_load(&current->next[0], memory_order_relaxed)
        );
        free(current);
        current = next;
    }
}

bool skiplist_insert(skiplist_t* list, int key, void* value) {
    skiplist_node_t* update[MAX_LEVEL];
    skiplist_node_t* current = list->head;
    
    // 找到每一层的插入位置
    for (int i = atomic_load(&list->level, memory_order_acquire) - 1; i >= 0; i--) {
        while (1) {
            tagged_ptr_t next_tagged = atomic_tagged_load(
                &current->next[i],
                memory_order_acquire
            );
            skiplist_node_t* next = get_ptr(next_tagged);
            
            if (!next || next->key > key) {
                update[i] = current;
                break;
            }
            if (next->key == key) {
                return false; // 键已存在
            }
            current = next;
        }
    }
    
    int new_level = random_level();
    skiplist_node_t* new_node = create_node(key, value, new_level);
    if (!new_node) return false;
    
    // 更新每一层的链接
    for (int i = 0; i < new_level; i++) {
        do {
            skiplist_node_t* next = get_ptr(
                atomic_tagged_load(&update[i]->next[i], memory_order_acquire)
            );
            new_node->next[i] = atomic_tagged_load(
                &update[i]->next[i],
                memory_order_relaxed
            );
            tagged_ptr_t new_next = make_tagged_ptr(
                new_node,
                get_tag(atomic_tagged_load(&update[i]->next[i],
                                         memory_order_relaxed)) + 1
            );
            
            if (atomic_tagged_compare_exchange_weak(
                    &update[i]->next[i],
                    &new_node->next[i],
                    new_next,
                    memory_order_release,
                    memory_order_relaxed)) {
                break;
            }
            cpu_pause();
        } while (1);
    }
    
    // 更新跳表层数
    while (atomic_load(&list->level, memory_order_acquire) < new_level) {
        atomic_compare_exchange_weak(
            &list->level,
            &(int){atomic_load(&list->level, memory_order_relaxed)},
            new_level,
            memory_order_release,
            memory_order_relaxed
        );
    }
    
    return true;
}

void* skiplist_find(skiplist_t* list, int key) {
    skiplist_node_t* current = list->head;
    
    for (int i = atomic_load(&list->level, memory_order_acquire) - 1; i >= 0; i--) {
        while (1) {
            tagged_ptr_t next_tagged = atomic_tagged_load(
                &current->next[i],
                memory_order_acquire
            );
            skiplist_node_t* next = get_ptr(next_tagged);
            
            if (!next || next->key > key) {
                break;
            }
            if (next->key == key) {
                return next->value;
            }
            current = next;
        }
    }
    
    return NULL;
}

bool skiplist_remove(skiplist_t* list, int key) {
    skiplist_node_t* update[MAX_LEVEL];
    skiplist_node_t* current = list->head;
    skiplist_node_t* target = NULL;
    
    // 找到每一层的删除位置
    for (int i = atomic_load(&list->level, memory_order_acquire) - 1; i >= 0; i--) {
        while (1) {
            tagged_ptr_t next_tagged = atomic_tagged_load(
                &current->next[i],
                memory_order_acquire
            );
            skiplist_node_t* next = get_ptr(next_tagged);
            
            if (!next || next->key > key) {
                update[i] = current;
                break;
            }
            if (next->key == key) {
                target = next;
            }
            current = next;
        }
    }
    
    if (!target) {
        return false;
    }
    
    // 从每一层删除节点
    for (int i = target->level - 1; i >= 0; i--) {
        do {
            tagged_ptr_t target_next = atomic_tagged_load(
                &target->next[i],
                memory_order_acquire
            );
            tagged_ptr_t new_next = make_tagged_ptr(
                get_ptr(target_next),
                get_tag(atomic_tagged_load(&update[i]->next[i],
                                         memory_order_relaxed)) + 1
            );
            
            if (atomic_tagged_compare_exchange_weak(
                    &update[i]->next[i],
                    &(tagged_ptr_t){target, get_tag(target_next)},
                    new_next,
                    memory_order_release,
                    memory_order_relaxed)) {
                break;
            }
            cpu_pause();
        } while (1);
    }
    
    free(target);
    return true;
}

size_t skiplist_size(skiplist_t* list) {
    size_t count = 0;
    skiplist_node_t* current = get_ptr(
        atomic_tagged_load(&list->head->next[0], memory_order_acquire)
    );
    
    while (current) {
        count++;
        current = get_ptr(
            atomic_tagged_load(&current->next[0], memory_order_acquire)
        );
    }
    
    return count;
} 